# 识别简单数学表达式验证码并登录
# 使用selenium模拟浏览器操作，使用ddddocr进行验证码识别
# pip install selenium ddddocr opencv-python requests pillow

import time
import requests
from selenium import webdriver
import cv2
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException
import ddddocr
from PIL import Image
import io
import re

class LoginBot:
    def __init__(self, login_url, username, password):
        self.login_url = login_url
        self.username = username
        self.password = password
        self.driver = webdriver.Chrome()  # 确保已安装ChromeDriver
        self.wait = WebDriverWait(self.driver, 10)
        # 初始化ddddocr
        self.ocr = ddddocr.DdddOcr()

    def login(self):
        """
        执行登录流程
        """
        try:
            # 打开登录页面
            self.driver.get(self.login_url)

            # 等待5秒后再开始执行登录流程
            print("页面已打开，等待20秒...")
            time.sleep(20)
            print("开始执行登录流程...")

            # 输入用户名
            username_input = self.wait.until(
                EC.presence_of_element_located((By.ID, "form_item_username"))
            )
            username_input.clear()
            username_input.send_keys(self.username)

            # 输入密码
            password_input = self.driver.find_element(By.ID, "form_item_password")
            password_input.clear()
            password_input.send_keys(self.password)

            # 检查是否有验证码
            captcha_input = None
            try:
                captcha_input = self.driver.find_element(By.ID, "form_item_code")
            except NoSuchElementException:
                print("未找到验证码输入框，继续登录流程")

            retry_count = 0
            max_retries = 100

            while retry_count < max_retries:
                if captcha_input:
                    # 处理验证码
                    captcha_code = self.recognize_captcha()
                    # print(f"已找到验证码输入框，开始处理验证码{captcha_code}")
                    if captcha_code:
                        captcha_input.clear()
                        captcha_input.send_keys(captcha_code)
                        print(f"输入验证码: {captcha_code}")
                    else:
                        # print("验证码识别失败")
                        if retry_count < max_retries - 1:
                            self.refresh_captcha()
                            retry_count += 1
                            continue
                        else:
                            raise Exception("验证码识别失败，已达到最大重试次数")

                # 点击登录按钮
                login_button = self.driver.find_element(By.XPATH, "//button[@type='submit']")
                login_button.click()

                # 等待页面响应
                time.sleep(2)

                # 检查是否提示验证码失效
                if self.check_captcha_expired():
                    # print("验证码已失效，刷新验证码并重新尝试")
                    self.refresh_captcha()
                    retry_count += 1
                    continue
                else:
                    # 检查是否登录失败（如密码错误、验证码错误等）
                    if self.check_login_failed():
                        # print("登录失败，刷新验证码并重新尝试")
                        # 清空验证码输入框
                        if captcha_input:
                            captcha_input.clear()
                        # 刷新验证码
                        self.refresh_captcha()
                        retry_count += 1
                        continue
                    else:
                        # 登录成功或失败但不是验证码问题
                        break

            # 检查是否登录成功
            if self.is_logged_in():
                print("登录成功")
                return True
            else:
                print("登录失败")
                return False

        except TimeoutException:
            print("页面加载超时")
            return False
        except Exception as e:
            print(f"登录过程中出现错误: {str(e)}")
            return False

    def recognize_captcha(self):
        """
        识别验证码图片
        使用dddocr进行验证码识别
        """
        try:
            # 等待验证码图片加载
            captcha_img = self.wait.until(
                EC.presence_of_element_located((By.ID, "getCaptcha"))
            )

            # 获取验证码图片的src或直接截图
            captcha_src = captcha_img.get_attribute("src")
            if captcha_src:
                # 如果是base64编码的图片
                if captcha_src.startswith("data:image"):
                    # 处理base64图片
                    import base64
                    header, encoded = captcha_src.split(",", 1)
                    image_data = base64.b64decode(encoded)
                    # 直接使用dddocr识别
                    result = self.ocr.classification(image_data)
                    print(f"原始识别结果: {result}")
                    return self.process_captcha_result(result)
                else:
                    # 从URL下载图片
                    response = requests.get(captcha_src)
                    # 使用dddocr识别
                    result = self.ocr.classification(response.content)
                    print(f"原始识别结果: {result}")
                    return self.process_captcha_result(result)
            else:
                # 直接截图验证码
                location = captcha_img.location
                size = captcha_img.size
                left = location['x']
                top = location['y']
                right = location['x'] + size['width']
                bottom = location['y'] + size['height']

                # 截取整个页面
                screenshot = self.driver.get_screenshot_as_png()
                full_image = Image.open(io.BytesIO(screenshot))

                # 裁剪验证码区域
                image = full_image.crop((left, top, right, bottom))
                
                # 转换为bytes
                img_bytes = io.BytesIO()
                image.save(img_bytes, format='PNG')
                img_bytes = img_bytes.getvalue()
                
                # 使用dddocr识别
                result = self.ocr.classification(img_bytes)
                print(f"原始识别结果: {result}")
                return self.process_captcha_result(result)

        except Exception as e:
            # print(f"验证码识别出错: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def process_captcha_result(self, text):
        """
        处理验证码识别结果
        """
        try:
            # print(f"原始识别结果: {text}")
            
            # 替换常见的错误识别字符
            # 将t替换为+号，将z替换为=号，将7替换为=号
            corrected_text = text.replace('t', '+').replace('z', '=').replace('T', '+').replace('Z', '=').replace('7', '=')
            
            # 针对"1+2=?"格式的特殊处理
            if '=' in corrected_text and '?' in corrected_text:
                # 提取等号前的表达式，过滤掉=和?部分
                expression = corrected_text.split('=')[0].strip()
                # print(f"提取到的表达式: {expression}")
                
                # 计算表达式结果
                result = self.safe_eval(expression)
                if result is not None:
                    # print(f"计算结果: {result}")
                    return result
            
            # 处理数学表达式
            # 先过滤掉可能的=和?字符
            cleaned_text = corrected_text.rstrip('=?').strip()
            if any(op in cleaned_text for op in ['+', '-', '*', '/']):
                # 查找表达式中的数字和运算符
                # 使用更宽松的正则表达式匹配数字（支持多位数）
                numbers = re.findall(r'\d+', cleaned_text)
                operators = re.findall(r'[+\-*/]', cleaned_text)
                
                # print(f"提取到的数字: {numbers}, 运算符: {operators}")
                
                if len(numbers) >= 2 and len(operators) >= 1:
                    # 构建表达式并计算结果
                    # 使用更智能的方式构建表达式
                    expression_parts = []
                    for i in range(max(len(numbers), len(operators))):
                        if i < len(numbers):
                            expression_parts.append(numbers[i])
                        if i < len(operators):
                            expression_parts.append(operators[i])
                    
                    expression = "".join(expression_parts)
                    # print(f"构建的表达式: {expression}")
                    
                    # 计算表达式结果
                    try:
                        # 安全计算表达式（避免使用eval）
                        result = self.safe_eval(expression)
                        if result is not None:
                            # print(f"计算表达式: {expression} = {result}")
                            return result
                    except Exception as e:
                        # print(f"表达式计算出错: {str(e)}")
                        # 尝试更简单的处理方式
                        if len(numbers) >= 2:
                            # 默认使用前两个数字和第一个运算符
                            try:
                                if operators[0] == '+':
                                    result = int(numbers[0]) + int(numbers[1])
                                elif operators[0] == '-':
                                    result = int(numbers[0]) - int(numbers[1])
                                elif operators[0] == '*':
                                    result = int(numbers[0]) * int(numbers[1])
                                elif operators[0] == '/':
                                    result = int(numbers[0]) // int(numbers[1]) if int(numbers[1]) != 0 else 0
                                # print(f"备用计算结果: {numbers[0]} {operators[0]} {numbers[1]} = {result}")
                                return result
                            except Exception as e2:
                                # print(f"备用计算也出错: {str(e2)}")
                                pass
                        
            # 如果不是数学表达式，返回清理后的文本
            if text:
                # 清理识别结果，只保留数字
                cleaned_text = re.sub(r'[^\d]', '', text.strip())
                if cleaned_text:
                    return cleaned_text
                
            return None

        except Exception as e:
            # print(f"处理验证码结果出错: {str(e)}")
            return None
    
    def recognize_math_question_captcha(self, image):
        """
        专门用于识别数学问题验证码（格式如：1+2=?）
        使用dddocr进行识别
        """
        try:
            # 转换为bytes
            img_bytes = io.BytesIO()
            image.save(img_bytes, format='PNG')
            img_bytes = img_bytes.getvalue()
            
            # 使用dddocr识别
            result = self.ocr.classification(img_bytes)
            print(f"原始识别结果: {result}")
            
            # 替换常见的错误识别字符
            # 将t替换为+号，将z替换为=号，将7替换为=号
            corrected_result = result.replace('t', '+').replace('z', '=').replace('T', '+').replace('Z', '=').replace('7', '=').replace('i', '/').replace('I', '/')
            
            # 处理识别结果
            if corrected_result and '=' in corrected_result and '?' in corrected_result:
                # 提取等号前的表达式，过滤掉=和?部分
                expression = corrected_result.split('=')[0].strip()
                # print(f"清理后的识别结果: {cleaned_text}")
                
                # print(f"计算前的表达式: {expression}")
                # 计算表达式结果
                result_value = self.safe_eval(expression)
                return result_value
            
            return None
            
        except Exception as e:
            # print(f"数学问题验证码识别出错: {str(e)}")
            return None
    
    def recognize_simple_math_captcha(self, image_gray):
        """
        专门用于识别简单数学验证码（加减法）的方法
        使用dddocr进行识别
        """
        try:
            # 转换为PIL图像
            image = Image.fromarray(image_gray)
            
            # 转换为bytes
            img_bytes = io.BytesIO()
            image.save(img_bytes, format='PNG')
            img_bytes = img_bytes.getvalue()
            
            # 使用dddocr识别
            result = self.ocr.classification(img_bytes)
            print(f"原始识别结果: {result}")
            
            # 替换常见的错误识别字符
            # 将t替换为+号，将z替换为=号，将7替换为=号
            corrected_result = result.replace('t', '+').replace('z', '=').replace('T', '+').replace('Z', '=').replace('7', '=')
            
            # 处理识别结果，过滤掉最后的=和?字符
            if corrected_result:
                # 先过滤掉可能的=和?字符
                cleaned_result = corrected_result.rstrip('=?').strip()
                # print(f"清理后的识别结果: {cleaned_result}")
                
                # 检查是否是有效的数学表达式
                if any(op in cleaned_result for op in ['+', '-']):
                    numbers = re.findall(r'\d+', cleaned_result)
                    operators = re.findall(r'[+\-]', cleaned_result)
                    
                    if len(numbers) >= 2:
                        # 构建表达式
                        expression = ""
                        for i in range(min(len(numbers), len(operators) + 1)):
                            expression += numbers[i]
                            if i < len(operators):
                                expression += operators[i]
                        
                        # print(f"计算前的表达式: {expression}")
                        # 计算结果
                        result_value = self.safe_eval(expression)
                        if result_value is not None:
                            return result_value
                elif len(re.findall(r'\d+', cleaned_result)) >= 2:
                    # 如果至少识别到两个数字
                    numbers = re.findall(r'\d+', cleaned_result)
                    expression = f"{numbers[0]}+{numbers[1]}"  # 默认使用加法
                    # print(f"重构的表达式: {expression}")
                    result_value = self.safe_eval(expression)
                    if result_value is not None:
                        return result_value
                        
            return None
        except Exception as e:
            # print(f"简单数学验证码识别出错: {str(e)}")
            return None

    def safe_eval(self, expression):
        """
        安全计算数学表达式
        """
        try:
            # 移除空格
            expression = expression.replace(" ", "")
            # 检查表达式是否只包含数字和运算符
            if not re.match(r'^[0-9+\-*/.]+$', expression):
                return None
            
            # 使用eval计算表达式，但限制操作
            allowed_chars = set('0123456789+-*/.() ')
            if all(c in allowed_chars for c in expression):
                # 替换可能的除零情况
                if '/0' in expression and not '/0.' in expression:
                    expression = expression.replace('/0', '/1')  # 避免除零错误
                
                result = eval(expression)
                # 如果结果是浮点数，检查是否为整数
                if isinstance(result, float):
                    if result.is_integer():
                        return int(result)
                    else:
                        # 四舍五入到整数
                        return round(result)
                return result
            return None
        except:
            return None

    def refresh_captcha(self):
        """
        刷新验证码
        """
        try:
            refresh_button = self.driver.find_element(By.ID, "refresh_captcha")
            refresh_button.click()
            time.sleep(1)  # 等待新验证码加载
        except NoSuchElementException:
            # 如果没有刷新按钮，尝试重新加载验证码图片
            try:
                captcha_img = self.driver.find_element(By.ID, "getCaptcha")
                captcha_img.click()
                time.sleep(1)
            except NoSuchElementException:
                # print("未找到验证码刷新机制")
                pass

    def check_captcha_expired(self):
        """
        检查是否提示验证码已失效
        """
        try:
            # 查找包含"验证码已失效"的元素
            error_message = self.driver.find_element(
                By.XPATH,
                "//*[contains(text(), '验证码已失效') or contains(text(), '验证码失效')]"
            )
            return error_message.is_displayed()
        except NoSuchElementException:
            return False

    def check_login_failed(self):
        """
        检查是否登录失败
        """
        try:
            # 查找包含登录失败信息的元素
            error_messages = self.driver.find_elements(
                By.XPATH,
                "//*[contains(text(), '密码错误') or contains(text(), '用户名错误') or contains(text(), '登录失败') or contains(text(), '验证码错误') or contains(@class, 'error') or contains(text(), '不正确')]"
            )
            
            # 检查是否有错误信息显示
            for element in error_messages:
                if element.is_displayed():
                    # print(f"检测到登录错误信息: {element.text}")
                    return True
            
            # 检查URL是否包含登录关键字（表示仍在登录页面）
            current_url = self.driver.current_url
            if "login" in current_url.lower():
                # 检查是否有表单输入框仍然可用（表示未跳转）
                try:
                    username_input = self.driver.find_element(By.ID, "form_item_username")
                    # 如果能找到用户名输入框且可用，可能表示登录未成功
                    if username_input.is_displayed() and username_input.is_enabled():
                        # 检查页面是否有错误提示（即使元素存在但可能不显示）
                        page_source = self.driver.page_source
                        error_keywords = ['密码错误', '用户名错误', '登录失败', '验证码错误', '不正确']
                        for keyword in error_keywords:
                            if keyword in page_source:
                                # print(f"在页面源码中检测到登录错误关键词: {keyword}")
                                return True
                except NoSuchElementException:
                    pass
            
            return False
        except Exception as e:
            # print(f"检查登录失败时出错: {str(e)}")
            return False

    def is_logged_in(self):
        """
        检查是否登录成功
        通过检查页面中接口回调是否返回成功来判断
        最多等待30秒
        """
        try:
            # 等待最多30秒检查登录结果
            end_time = time.time() + 30
            while time.time() < end_time:
                # 检查是否存在登录成功的标识
                try:
                    # 检查是否有登录成功的特定元素或文本
                    success_elements = self.driver.find_elements(
                        By.XPATH,
                        "//*[contains(text(), '登录成功') or contains(text(), '欢迎') or contains(text(), '退出') or contains(@class, 'user-info') or contains(@class, 'welcome')]"
                    )
                    
                    # 检查页面URL是否变化（通常登录成功后会跳转到其他页面）
                    current_url = self.driver.current_url
                    if "login" not in current_url.lower():
                        # print("检测到URL变化，登录可能成功")
                        return True
                    
                    # 检查是否有明显的成功标识元素
                    if success_elements:
                        for element in success_elements:
                            if element.is_displayed():
                                # print("检测到登录成功标识元素")
                                return True
                    
                    # 检查是否有特定的响应数据（例如通过JavaScript变量）
                    try:
                        # 检查是否存在表示登录成功的JavaScript变量或状态
                        login_status = self.driver.execute_script("return window.loginStatus || window.isLoggedIn || null;")
                        if login_status is True:
                            # print("检测到JavaScript登录成功状态")
                            return True
                    except:
                        pass  # 忽略JavaScript执行错误
                    
                    # 检查是否有错误信息，如果有则直接返回失败
                    error_elements = self.driver.find_elements(
                        By.XPATH,
                        "//*[contains(text(), '密码错误') or contains(text(), '用户名错误') or contains(text(), '登录失败') or contains(text(), '验证码错误') or contains(@class, 'error')]"
                    )
                    if error_elements:
                        for element in error_elements:
                            if element.is_displayed():
                                # print("检测到登录错误信息，登录失败")
                                return False
                    
                except NoSuchElementException:
                    pass
                
                # 等待一小段时间再检查
                time.sleep(1)
            
            # 超时后仍未检测到明确的成功或失败标识
            # print("等待30秒后仍未检测到明确的登录结果")
            # 再次检查URL
            current_url = self.driver.current_url
            if "login" not in current_url.lower():
                return True
            return False
            
        except Exception as e:
            # print(f"检查登录状态时出错: {str(e)}")
            # 出错时检查URL
            try:
                current_url = self.driver.current_url
                if "login" not in current_url.lower():
                    return True
            except:
                pass
            return False

    def close(self):
        """
        关闭浏览器
        """
        self.driver.quit()


# 使用示例
if __name__ == "__main__":
    # 配置登录信息
    LOGIN_URL = "http://localhost:8090/login"  # 替换为实际的登录页面URL
    USERNAME = "admin"  # 替换为实际用户名
    PASSWORD = "super-admin@2025!"  # 替换为实际密码

    # 创建登录机器人实例
    bot = LoginBot(LOGIN_URL, USERNAME, PASSWORD)

    try:
        # 执行登录
        success = bot.login()
        if success:
            print("登录流程完成")
            # 在这里可以继续执行其他操作
            time.sleep(5)  # 等待5秒再关闭
        else:
            print("登录失败")
    finally:
        # 关闭浏览器
        bot.close()